/*
 *******************************************************************
 * Author: xuelei
 * Create Date: 2021-09-09 14:12
 * Description: create file
 *******************************************************************
 */
package com.xuelei.cloud.bc.crypto.asymmetric;

import com.xuelei.cloud.bc.Checker;
import com.xuelei.cloud.bc.Provider;
import com.xuelei.cloud.bc.crypto.exception.CryptoException;
import org.bouncycastle.asn1.gm.GMNamedCurves;
import org.bouncycastle.asn1.x9.X9ECParameters;
import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.InvalidCipherTextException;
import org.bouncycastle.crypto.engines.SM2Engine;
import org.bouncycastle.crypto.params.ECDomainParameters;
import org.bouncycastle.crypto.params.ECPrivateKeyParameters;
import org.bouncycastle.crypto.params.ECPublicKeyParameters;
import org.bouncycastle.crypto.params.ParametersWithRandom;
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPrivateKey;
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPublicKey;
import org.bouncycastle.math.ec.ECCurve;
import org.bouncycastle.math.ec.ECPoint;
import org.bouncycastle.util.BigIntegers;
import org.bouncycastle.util.encoders.Hex;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.math.BigInteger;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.spec.ECGenParameterSpec;

/**
 * 基于SM2算法的加密解密工具类
 * @author xuelei
 * @date 2021-09-09
 */
public class SM2Util {
    
    private static final Logger log = LoggerFactory.getLogger(SM2Util.class);
    
    private static final String SM2_EC_CURVE = "sm2p256v1";
    
    private static SecureRandom secureRandom;
    
    private static final String SECURITY_RANDOM_ALGORITHM = "SHA1PRNG";
    
    private static final boolean ENCRYPT = true;
    
    private static final boolean DECRYPT = false;
    
    private static final Charset DEFAULT_CHARSET = StandardCharsets.UTF_8;
    
    private static final byte Ox_02 = 0x02;
    
    private static final byte Ox_03 = 0x03;
    
    private static final byte Ox_04 = 0x04;
    
    private static final String DEFAULT_SM2_MODE = SM2Mode.C1C3C2;
    
    static {
        try {
            secureRandom = SecureRandom.getInstance(SECURITY_RANDOM_ALGORITHM);
            secureRandom.setSeed(secureRandom.generateSeed(8));
        } catch (Exception e) {
            log.error("Exception: {}, message: {}", e.getClass(), e.getMessage());
        }
    }
    
    /**
     * 获取SM2公钥私钥对
     * @return 公钥私钥对, 以Base16形式对密钥进行编码
     * @throws CryptoException
     */
    public static SM2KeyPair createKeyPair() throws CryptoException {
        KeyPair keyPair = SM2Util.createKey();
        BCECPrivateKey privateKey = (BCECPrivateKey) keyPair.getPrivate();
        BCECPublicKey publicKey = (BCECPublicKey) keyPair.getPublic();
        SM2KeyPair sm2KeyPair = new SM2KeyPair();
        byte[] privateKeyBlock = BigIntegers.asUnsignedByteArray(32,privateKey.getD());
        byte[] publicKeyXBlock = BigIntegers.asUnsignedByteArray(32, publicKey.getQ().getRawXCoord().toBigInteger());
        byte[] publicKeyYBlock = BigIntegers.asUnsignedByteArray(32, publicKey.getQ().getRawYCoord().toBigInteger());
        sm2KeyPair.setPrivateKey(Hex.toHexString(privateKeyBlock));
        sm2KeyPair.setPublicKeyX(Hex.toHexString(publicKeyXBlock));
        sm2KeyPair.setPublicKeyY(Hex.toHexString(publicKeyYBlock));
        return sm2KeyPair;
    }
    
    /**
     * 加密字符串，以C1C3C2模式组织密文
     * @param plainText 明文字符串
     * @param publicKeyXBase16 Base16编码的公钥X值
     * @param publicKeyYBase16 Base16编码的公钥X值
     * @return Base16编码的密文字符串
     * @throws CryptoException
     */
    public static String encrypt(String plainText, String publicKeyXBase16, String publicKeyYBase16) throws CryptoException {
        return SM2Util.encrypt(plainText, publicKeyXBase16, publicKeyYBase16, DEFAULT_SM2_MODE);
    }
    
    /**
     * 加密字符串
     * @param plainText 明文字符串
     * @param publicKeyXBase16 Base16编码的公钥X值
     * @param publicKeyYBase16 Base16编码的公钥X值
     * @param mode 密文组织模式，参见{@link SM2Mode}
     * @return Base16编码的密文字符串
     * @throws CryptoException
     */
    public static String encrypt(String plainText, String publicKeyXBase16, String publicKeyYBase16, String mode) throws CryptoException {
        Checker.checkText(plainText, "plain text is empty");
        byte[] plainBlock = plainText.getBytes(DEFAULT_CHARSET);
        byte[] cipherBlock = SM2Util.encrypt(plainBlock, publicKeyXBase16, publicKeyYBase16, mode);
        return Hex.toHexString(cipherBlock);
    }
    
    /**
     * 加密数据块，以C1C3C2模式组织密文
     * @param plainBlock 明文数据块
     * @param publicKeyXBase16 Base16编码的公钥X值
     * @param publicKeyYBase16 Base16编码的公钥X值
     * @return 加密后的数据块
     * @throws CryptoException
     */
    public static byte[] encrypt(byte[] plainBlock, String publicKeyXBase16, String publicKeyYBase16) throws CryptoException {
        return SM2Util.encrypt(plainBlock, publicKeyXBase16, publicKeyYBase16, DEFAULT_SM2_MODE);
    }
    
    /**
     * 加密数据块
     * @param plainBlock 明文数据块
     * @param publicKeyXBase16 Base16编码的公钥X值
     * @param publicKeyYBase16 Base16编码的公钥X值
     * @param mode 密文组织模式，参见{@link SM2Mode}
     * @return 加密后的数据块
     * @throws CryptoException
     */
    public static byte[] encrypt(byte[] plainBlock, String publicKeyXBase16, String publicKeyYBase16, String mode) throws CryptoException {
        Checker.checkBlock(plainBlock, "plain block is empty");
        byte[] cipherBlock = null;
        CipherParameters publicKeyParameter = SM2Util.buildECPublicKeyParameter(publicKeyXBase16, publicKeyYBase16);
        SM2Engine sm2Engine = new SM2Engine(SM2Engine.Mode.valueOf(mode));
        sm2Engine.init(ENCRYPT, publicKeyParameter);
        try {
            cipherBlock = sm2Engine.processBlock(plainBlock, 0, plainBlock.length);
        } catch (InvalidCipherTextException e) {
            log.error("Exception: {}, message: {}", e.getClass(), e.getMessage());
            e.printStackTrace();
            throw new CryptoException(e.getMessage());
        }
        return cipherBlock;
    }
    
    /**
     * 解密以C1C3C2模式组织的密文数据块
     * @param cipherText 以C1C3C2模式组织的密文字符串
     * @param privateKeyBase16 Base16编码的私钥
     * @return 解密后的明文字符串
     * @throws CryptoException
     */
    public static String decrypt(String cipherText, String privateKeyBase16) throws CryptoException {
        return SM2Util.decrypt(cipherText, privateKeyBase16, DEFAULT_SM2_MODE);
    }
    
    /**
     * 解密密文字符串
     * @param cipherText 密文字符串
     * @param privateKeyBase16 Base16编码的私钥
     * @param mode 密文组织模式，参见{@link SM2Mode}
     * @return 解密后的明文字符串
     * @throws CryptoException
     */
    public static String decrypt(String cipherText, String privateKeyBase16, String mode) throws CryptoException {
        Checker.checkText(cipherText, "cipher text is empty");
        byte[] cipherBlock = Hex.decode(cipherText);
        byte[] plainBlock = SM2Util.decrypt(cipherBlock, privateKeyBase16, mode);
        return new String(plainBlock);
    }
    
    /**
     * 解密以C1C3C2模式组织的密文数据块
     * @param cipherBlock 以C1C3C2模式组织的密文数据块
     * @param privateKeyBase16 Base16编码的私钥
     * @return 解密后的明文数据块
     * @throws CryptoException
     */
    public static byte[] decrypt(byte[] cipherBlock, String privateKeyBase16) throws CryptoException {
        return SM2Util.decrypt(cipherBlock, privateKeyBase16, DEFAULT_SM2_MODE);
    }
    
    /**
     * 解密密文数据块
     * @param cipherBlock 密文数据块
     * @param privateKeyBase16 Base16编码的私钥
     * @param mode 密文组织模式，参见{@link SM2Mode}
     * @return 解密后的明文数据块
     * @throws CryptoException
     */
    public static byte[] decrypt(byte[] cipherBlock, String privateKeyBase16, String mode) throws CryptoException {
        Checker.checkBlock(cipherBlock, "cipher block is empty");
        byte[] plainBlock = null;
        cipherBlock = standardize(cipherBlock);
        CipherParameters privateKeyParameter = SM2Util.buildECPrivateKeyParameter(privateKeyBase16);
        SM2Engine sm2Engine = new SM2Engine(SM2Engine.Mode.valueOf(mode));
        sm2Engine.init(DECRYPT, privateKeyParameter);
        try {
            plainBlock = sm2Engine.processBlock(cipherBlock, 0, cipherBlock.length);
        } catch (InvalidCipherTextException e) {
            log.error("Exception: {}, message: {}", e.getClass(), e.getMessage());
            throw new CryptoException(e.getMessage());
        }
        return plainBlock;
    }
    
    private static KeyPair createKey() throws CryptoException {
        KeyPair keyPair = null;
        try {
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("EC", Provider.BouncyCastle);
            ECGenParameterSpec ecGenParameterSpec = new ECGenParameterSpec(SM2_EC_CURVE);
            keyPairGenerator.initialize(ecGenParameterSpec);
            keyPair = keyPairGenerator.generateKeyPair();
        } catch (NoSuchAlgorithmException | InvalidAlgorithmParameterException e) {
            log.error("Exception: {}, message: {}", e.getClass(), e.getMessage());
            throw new CryptoException(e.getMessage());
        }
        return keyPair;
    }
    
    private static CipherParameters buildECPrivateKeyParameter(String privateKeyBase16) {
        ECDomainParameters ecDomainParameters = SM2Util.createECDomainParameters();
        BigInteger privateKey = BigIntegers.fromUnsignedByteArray(Hex.decode(privateKeyBase16));
        return new ECPrivateKeyParameters(privateKey, ecDomainParameters);
    }
    
    private static CipherParameters buildECPublicKeyParameter(String publicKeyXBase16, String publicKeyYBase16) {
        ECDomainParameters ecDomainParameters = SM2Util.createECDomainParameters();
        ECCurve curve = ecDomainParameters.getCurve();
        BigInteger x = BigIntegers.fromUnsignedByteArray(Hex.decode(publicKeyXBase16));
        BigInteger y = BigIntegers.fromUnsignedByteArray(Hex.decode(publicKeyYBase16));
        ECPoint point = curve.createPoint(x, y);
        ECPublicKeyParameters publicKeyParameters = new ECPublicKeyParameters(point, ecDomainParameters);
        return new ParametersWithRandom(publicKeyParameters);
    }
    
    private static ECDomainParameters createECDomainParameters() {
        X9ECParameters x9ECParameters = GMNamedCurves.getByName(SM2_EC_CURVE);
        return new ECDomainParameters(x9ECParameters);
    }
    
    private static byte[] standardize(final byte[] cipherBlock) {
        if (cipherBlock[0] == Ox_02 || cipherBlock[0] == Ox_03 || cipherBlock[0] == Ox_04) {
            return cipherBlock;
        } else {
            byte[] standardizeCipherBlock = new byte[cipherBlock.length + 1];
            standardizeCipherBlock[0] = Ox_04;
            System.arraycopy(cipherBlock, 0, standardizeCipherBlock, 1, cipherBlock.length);
            return standardizeCipherBlock;
        }
    }
    
    public static class SM2KeyPair {
        
        private String publicKeyX;
        
        private String publicKeyY;
        
        private String privateKey;
    
        public String getPublicKeyX() {
            return publicKeyX;
        }
    
        private void setPublicKeyX(String publicKeyX) {
            this.publicKeyX = publicKeyX;
        }
    
        public String getPublicKeyY() {
            return publicKeyY;
        }
    
        private void setPublicKeyY(String publicKeyY) {
            this.publicKeyY = publicKeyY;
        }
    
        public String getPrivateKey() {
            return privateKey;
        }
    
        private void setPrivateKey(String privateKey) {
            this.privateKey = privateKey;
        }
    }
}
